package Algorithm;

import Structure.Edge;
import Structure.Graph;
import Structure.GraphGenerator;
import Structure.Node;

import java.util.Collection;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.PriorityQueue;
import java.util.Set;
/**
 * @Author: ylc
 * @Description:Kruskal算法
 * @Date Created in 2022-06-07 19:02
 */
public class Kruskal {
    public static class UnionFind {
        private HashMap<Node, Node> fatherMap;
        private HashMap<Node, Integer> rankMap;

        public UnionFind() {
            fatherMap = new HashMap<Node, Node>();
            rankMap = new HashMap<Node, Integer>();
        }

        private Node findFather(Node n) {
            Node father = fatherMap.get(n);
            if (father != n) {
                father = findFather(father);
            }
            fatherMap.put(n, father);
            return father;
        }

        public void makeSets(Collection<Node> nodes) {
            fatherMap.clear();
            rankMap.clear();
            for (Node node : nodes) {
                fatherMap.put(node, node);
                rankMap.put(node, 1);
            }
        }

        public boolean isSameSet(Node a, Node b) {
            return findFather(a) == findFather(b);
        }

        public void union(Node a, Node b) {
            if (a == null || b == null) {
                return;
            }
            Node aFather = findFather(a);
            Node bFather = findFather(b);
            if (aFather != bFather) {
                int aFrank = rankMap.get(aFather);
                int bFrank = rankMap.get(bFather);
                if (aFrank <= bFrank) {
                    fatherMap.put(aFather, bFather);
                    rankMap.put(bFather, aFrank + bFrank);
                } else {
                    fatherMap.put(bFather, aFather);
                    rankMap.put(aFather, aFrank + bFrank);
                }
            }
        }
    }

    public static class EdgeComparator implements Comparator<Edge> {

        @Override
        public int compare(Edge o1, Edge o2) {
            return o1.weight - o2.weight;
        }

    }

    /**
     * 克鲁斯卡尔最小生成树算法
     * @param graph 无向图
     * @return
     */
    public static Set<Edge> kruskalMST(Graph graph) {
        UnionFind unionFind = new UnionFind();
        unionFind.makeSets(graph.nodes.values());
        PriorityQueue<Edge> priorityQueue = new PriorityQueue<>(new EdgeComparator());
        for (Edge edge : graph.edges) {
            priorityQueue.add(edge);
        }
        Set<Edge> result = new HashSet<>();
        while (!priorityQueue.isEmpty()) {
            Edge edge = priorityQueue.poll();
            if (!unionFind.isSameSet(edge.from, edge.to)) {
                result.add(edge);
                unionFind.union(edge.from, edge.to);
            }
        }
        return result;
    }
    public static void main(String[] args) {
        Graph graph = create();
        Set<Edge> edges = kruskalMST(graph);
        for (Edge edge :
                edges) {
            System.out.print(edge.weight + " ");
        }
    }
    public static Graph create(){
        int[][] arr = {
                {6,1,2},{6,2,1},
                {1,1,3},{1,3,1},
                {5,1,4},{5,4,1},
                {5,2,3},{5,3,2},
                {5,3,4},{5,4,3},
                {3,2,5},{3,5,2},
                {6,3,5},{6,5,3},
                {4,3,6},{4,6,3},
                {2,4,6},{2,6,4},
                {6,5,6},{6,6,5},
        };

        return GraphGenerator.generateGraph(arr);
    }
}
